[ROCm][MoE] Pad the AITER MoE intermediate size at allocation time, and round the expert-group count to a kernel that exists - #55368

Purpose

Two independent startup failures on the ROCm AITER MoE path, both caused by shapes AITER has no compiled kernel for, both focusing on IFM/K2-Horizon-375B-A23B model support.

1. Unaligned per-partition intermediate size - unquantized path only, fp8 AITER alignment is unchanged

AITER's CK 2stages MoE kernel dispatches on inter_dim <= 192: below it both stages use 64-wide tiles, above it at least one stage uses a 128-wide tile. CK's IsSupportedArgument rejects an intermediate not divisible by that width, so moe_intermediate_size = 1792 at TP=8 gives 224 and fails with device_gemm ... does not support this GEMM problem.

Fix: round the intermediate up in maybe_roundup_sizes, the existing hook for this (already overridden by five other quant methods). It runs before create_weights, so the padded size reaches allocation and the loader — which derives shard offsets from the checkpoint, not the parameter shape — fills only the real rows.

Alignment rule: 64 if inter_dim <= 192 else 128, mirroring AITER's own threshold rather than always using 128. A flat 128 mispredicts 2 of the 19 sizes measured against the real kernel (64 and 192 both pass despite not being 128-aligned) — it is a wrong model that happens to be safe. It also costs: inter_dim is a tuned-config lookup key, and across the 148 shipped AITER config files (11279 rows) the threshold rule pads 0 tuned rows while a flat 128 would move 207 off theirs onto the heuristic fallback. Rounding up never crosses the boundary (96->128, 160->192, 192->192 all stay <= 192), so the alignment picked is always the one the dispatcher then applies. Dtype-independent: CK derives both tile widths from sizeof(A0DataType), identical for fp16 and bf16.

Zero-init. The loader narrows to the checkpoint extent and never writes the tail, so w13/w2 are allocated with torch.zeros when either dim is padded — unpadded layers keep torch.empty and pay nothing. Pad lanes are then bit-exactly inert: silu(0) * 0 = 0, and the zero columns in w2 contribute exactly zero to stage 2. With torch.empty the tail holds recycled device memory, which is live weight to the kernel; whether that hurts depends on load-time allocation order, so it presents as sporadic implausible tokens with no error. mxfp4 and quark nvfp4 already allocate zeroed; #55251 is doing the same for fp8.

Scope. Gated on UnquantizedMoeBackend.AITER, only selected on ROCm with AITER MoE enabled; every other unquantized backend falls through the base hook untouched. The quantized AITER paths have their own overrides with different constants and are not touched — fp8 needs 256.

Known interaction. Since intermediate_size_per_partition now differs from ..._unpadded, experts/rocm_aiter_moe.py computes a non-zero intermediate_pad where it previously got 0. Traced: with QuantType.No + bf16 the dispatch falls past every cktile/flydsl branch (all require q_type == per_1x32) to the plain ck_moe_stage1/ck_moe_stage2_fwd partials, which take no pad argument, and intermediate_pad only enters the tuned-config key when config_file is not None, which vLLM never sets. So it is computed and discarded, and the kernel runs the full padded GEMM including the zero rows — up to ~14% extra compute on the MoE GEMMs (448->512 at TP=4, 224->256 at TP=8).

2. AITER biased_grouped_topk expert-group count

_aiter_get_num_expert_group ceil-divides num_experts by the 32-per-group limit then walks to the next divisor, which can produce a count with no compiled kernel: AITER instantiates biased_grouped_topk only for NUM_GRP in {1,2,4,8}, so 96 experts gives 3 and fails to launch.

Grouping is a no-op here (topk_group == num_expert_group), so any divisor within the limit routes identically — the constraint is purely which kernel exists. Round to the largest supported count that divides num_experts and respects the limit. When none fits (320 -> 10, 33 -> 3) the naive value is kept deliberately: it is large enough that the topk >= num_expert_group guard at the call site fails and routing falls back to the generic path rather than reaching AITER.

Test Plan

Both files are CPU-only and gated on current_platform.is_rocm(), so they run in any ROCm CI job without a working AITER runtime.

tests/kernels/moe/test_rocm_aiter_moe.py — new "Weight alignment" section appended to the existing file rather than a standalone module, following that file's conventions (function-local imports, default_vllm_config, rocm_ naming):

• test_aiter_moe_alignment_follows_threshold — the rule, as a table across the boundary.
• test_aiter_moe_padded_size_stays_in_its_dispatch_branch — self-consistency: (intermediate <= 192) == (padded <= 192) and alignment(padded) == alignment(intermediate).
• test_aiter_moe_roundup_pads_intermediate — the hook end-to-end: 96->128, 192->192, 224->256, 448->512, 512->512.
• test_aiter_moe_roundup_is_not_applied_to_other_backends — TRITON and FLASHINFER_CUTLASS keep 224.
• test_aiter_moe_padded_weights_are_zero_initialized — pad lanes are zero, not allocator garbage.
• test_aiter_moe_aligned_weights_are_zero_initialized — zero-init is unconditional, so it holds for an aligned size too.
• test_aiter_moe_bias_matches_padded_weight — w13_bias.shape[1] == w13_weight.shape[1] with has_bias=True; the defect the post-load approach had.
• test_aiter_moe_padding_is_numerically_inert — reproduces the loader's gate/up placement into a zeroed padded parameter, asserts the MoE output is bit-identical to the unpadded reference.

tests/kernels/moe/test_rocm_aiter_num_expert_group.py (new) — five tests over _aiter_get_num_expert_group: the two invariants the router asserts unconditionally, that a supported count is chosen whenever one fits, a table of known expert counts, and that the fallback is taken deliberately when none fits.

Test Result

gfx950 (MI355X), torch 2.11.0+gitd0c8b1f, on 29af8bd672.

pytest tests/kernels/moe/test_rocm_aiter_num_expert_group.py -q   49 passed, 2 skipped
pytest tests/kernels/moe/test_rocm_aiter_moe.py -q                 1 failed, 63 passed
both files                                                         1 failed, 112 passed, 2 skipped  (72s)

The one failure, test_aiter_fused_moe_mi3xx_fp8_accuracy, is pre-existing and unrelated — verified by running it alone against an unmodified checkout of the parent commit in a separate worktree, same box and image, where it fails identically.